import numpy as np
from skimage.io import imread
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import pickle
import time

print("开始处理...")

# 加载图像
print("\n正在读取图像...")
X_img = imread(r"C:\Users\29650\Desktop\yanashi\Sandstone_1.tif")
y_img = imread(r"C:\Users\29650\Desktop\yanashi\Sandstone_1_segment.tif")

print(f"图像形状: {X_img.shape}")
print(f"标签形状: {y_img.shape}")

# 如果y_img是三通道的，我们只取第一个通道或将其转换为单通道
if len(y_img.shape) == 3:
    y_img = y_img[:,:,0]  # 只取第一个通道
    print("将标签图像转换为单通道")

# 确保两个图像具有相同的尺寸
min_height = min(X_img.shape[0], y_img.shape[0])
min_width = min(X_img.shape[1], y_img.shape[1])

X_img = X_img[:min_height, :min_width]
y_img = y_img[:min_height, :min_width]

# 将图像重塑为二维数组
X = X_img.reshape(X_img.shape[0] * X_img.shape[1], -1)
y = y_img.flatten()

print("\n完成从砂岩截面图1及其对应分区中获取X和y")
print(f"处理后的X形状: {X.shape}")
print(f"处理后的y形状: {y.shape}")

# 将数据分为训练集和测试集
print("\n正在划分训练集和测试集...")
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
print("完成train_test_split")
print(f"训练集大小: {X_train.shape}")
print(f"测试集大小: {X_test.shape}")

# 训练随机森林分类器
print("\n开始训练随机森林模型...")
start_time = time.time()
clf = RandomForestClassifier(n_estimators=100, random_state=42)
clf.fit(X_train, y_train)
training_time = time.time() - start_time
print(f"完成随机森林模型clf的训练 (用时: {training_time:.2f}秒)")

# 评估模型
accuracy = clf.score(X_test, y_test)
print(f"准确率: {accuracy:.4f}")

# 保存训练好的模型
print("\n正在保存模型...")
with open('clf.pkl', 'wb') as f:
    pickle.dump(clf, f)
print("已保存随机森林模型clf到硬盘")

print("\n所有处理完成!")